import math
import logging
import pandas as pd
import backtrader as bt
import psycopg2.extras
from pyLibs import units, Exceptions, TdxSource, StockIndicator

logging.basicConfig(format="%(asctime)s %(filename)s(line:%(lineno)d) [%(levelname)s] : %(message)s",
                    datefmt="%Y-%M-%d %H:%M:%S", level=logging.DEBUG)


class TradeStrategy(bt.Strategy):
    params = (
        ('printlog', True),
        ('smoothing_period', 5),
        ('stack_len', 3),
    )

    """提供记录功能"""

    def log(self, txt, dt=None, doprint=False):
        ''' Logging function fot this strategy'''
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.datetime(0)
            print('%s: %s' % (dt.isoformat(sep=" "), txt))

    '''
    放量突破布林带中轨后进行平台整理，再向上的策略
        1、先标记放量突破布林带中轨
        2、中轨逐步升高
        3、长期EMA均线逐步升高
        4、突破中轨后记录最高价
        5、突破中轨后记录最低价
        6、突破中轨后至少有3根K线的收盘价在最高价和最低价范围内
        7、当向上突破最高价，收盘价高于最高价的0.5%，发出买入信号
        8、当价格跌破突破布林带中轨时的价位时，取消监控
    仓位管理：
        1、以突破BAR的收盘价挂限价单买多50%的仓位
        2、以整理平台的上沿价格挂限价单买多50%的仓位
    止盈止损：
        1、初始止损价为整理平台下方2个报价单位
        2、买入后再创新高，改用1.5倍ATR进行第一移动止损，止损50%仓位
        3、买入后再创新高，发生回撤，回撤的低点高于前高时，回撤低点为第二移动止损，否则以前高为第二移动止损，完全止损。
    modalitySign 的列名意义：
        stock_code                      股票代码
        datetime                        日期时间
        open                            开盘价
        close                           收盘价
        high                            最高价
        low                             最低价
        ema_s                           短期EMA
        ema_m                           中期EMA
        ema_l                           长期EMA
        atr                             ATR
        close_cross_mid                 向上突破中轨时的价格
        mid_ema_position                中轨与EMA相对位置（mid-ema）
        bollinger_mid                   布林带中轨价格
        platform_high                   区间阶段的高点
        platform_low                    区间阶段的低点
        is_cross_mid                    是否穿越布林带中轨
        is_vol_release                  是否放量
        is_vol_multiplier               是否倍量
        is_mid_increase                 中轨是否逐步抬高
        is_ema_increase                 长期EMA是否逐步抬高
        mid_cross_ema_m                 布林带中轨是否金叉中期EMA
        mid_cross_ema_l                 布林带中轨是否金叉长期EMA
        break_high_enter                是否发出向上突破平台上沿的做多信号
        break_low_enter                 是否发出向下突破平台下沿的做空信号
        initial_loss                    初始止损价位
        moving_lost                     移动止损价位
        is_trigger_initial              是否触发初始止损价位
        is_trigger_moving               是否触发移动止损价位
        status                          状态：Watch（观察）、GiveUp（放弃）、Signal（发出信号）、Enter（介入）、Wait（等待）、InitialLoss（触发初始止损）、MovingLoss（触发移动止损）
    '''

    def __init__(self):
        self.datasLength = len(self.getdatanames())
        self.watchSign = self.orderSign = False
        self.modality = None
        self.modalityList = list()
        self.columns = ['stock_code', 'datetime', 'open', 'close', 'high', 'low', 'ema_m', 'ema_l', 'bollinger_mid',
                        'mid_ema_position', 'platform_high', 'platform_low', 'atr', 'close_cross_mid', 'initial_loss',
                        'moving_lost', 'is_cross_mid', 'is_vol_release', 'is_vol_multiplier', 'is_mid_increase',
                        'is_ema_increase', 'mid_cross_ema_m', 'mid_cross_ema_l', 'break_high_enter', 'break_low_enter',
                        'is_trigger_initial', 'is_trigger_moving', 'status']
        # self.modalityDict = dict.fromkeys(tuple(self.columns))
        self.modalitySign = pd.DataFrame(columns=self.columns)
        self.indicators = list()
        for i in range(0, self.datasLength):
            bollinger = bt.indicators.BollingerBands(self.datas[i], period=20, devfactor=2.0)
            ema_m = bt.indicators.ExponentialMovingAverage(self.datas[i], period=55)
            ema_l = bt.indicators.ExponentialMovingAverage(self.datas[i], period=144)
            macd = bt.indicators.MACDHisto(self.datas[i], period_me1=12, period_me2=26, period_signal=9)
            self.indicators.append({
                'vol_s': bt.indicators.SmoothedMovingAverage(self.datas[i].volume, period=13),
                'vol_l': bt.indicators.SmoothedMovingAverage(self.datas[i].volume, period=89),
                'ema_m': ema_m,
                'ema_l': ema_l,
                'mid': bollinger.mid,
                'top': bollinger.top,
                'bot': bollinger.bot,
                'close_cross_mid': bt.indicators.CrossOver(self.datas[i].close, bollinger.mid),
                'mid_cross_emaM': bt.indicators.CrossOver(bollinger.mid, ema_m),
                'mid_cross_emaL': bt.indicators.CrossOver(bollinger.mid, ema_l),
                'diff': macd.macd,
                'dea': macd.signal,
                'macd': macd.histo,
                'atr': bt.indicators.AverageTrueRange(self.datas[i], period=14),
            })
        self.marketposition = 0  # 保存交易状态
        self.buyprice = None
        self.sellprice = None
        self.loss = 0
        # 信号栈长短
        self.stack = [0] * self.params.stack_len

    def next(self):
        # print('【%s】%s，收盘价：￥%.2f' % (self.datas[0]._name, self.datas[0].datetime.datetime(0), self.datas[0].close[0]))
        modality = dict.fromkeys(tuple(self.columns))
        for i in range(0, self.datasLength):
            stock = self.datas[i]
            indicators = self.indicators[i]
            is_vol_release = True if stock.volume[0] > indicators['vol_l'][0] else False
            modality['stock_code'] = stock._name
            modality['datetime'] = stock.datetime.datetime(0)
            modality['open'] = round(stock.open[0], 4)
            modality['close'] = round(stock.close[0], 4)
            modality['high'] = round(stock.high[0], 4)
            modality['low'] = round(stock.low[0], 4)
            modality['ema_m'] = round(indicators['ema_m'][0], 4)
            modality['ema_l'] = round(indicators['ema_l'][0], 4)
            modality['atr'] = round(indicators['atr'][0], 4)
            modality['bollinger_mid'] = round(indicators['mid'][0], 4)
            modality['mid_ema_position'] = round(indicators['mid'][0] - indicators['ema_m'][0], 4)
            modality['is_vol_release'] = is_vol_release
            modality['is_vol_multiplier'] = True if stock.volume[0] > indicators['vol_l'][0] * 2 else False
            modality['is_mid_increase'] = True if indicators['mid'][0] > indicators['mid'][-1] else False
            modality['is_ema_increase'] = True if indicators['ema_l'][0] > indicators['ema_l'][-1] else False
            modality['mid_cross_ema_m'] = True if indicators['mid_cross_emaM'][0] else False
            modality['mid_cross_ema_l'] = True if indicators['mid_cross_emaL'][0] else False
            if self.orderSign is not True:
                # 尚未创建订单
                if self.watchSign is not True:
                    # 尚未进入关注状态
                    is_cross_mid = True if indicators['close_cross_mid'][0] == 1 else False
                    if is_cross_mid:
                        # 放量向上突破布林带中轨
                        modality['close_cross_mid'] = indicators['mid'][0]
                        modality['platform_high'] = stock.high[0]
                        modality['is_cross_mid'] = is_cross_mid
                        modality['status'] = 'Watch'
                        self.modalityList.append(modality)
                        self.watchSign = True
                else:
                    # 已经进入关注状态，接下来判断是否创出新高
                    modality['close_cross_mid'] = self.modalityList[-1:][0]['close_cross_mid']
                    modality['platform_high'] = self.modalityList[-1:][0]['platform_high']
                    modality['platform_low'] = self.modalityList[-1:][0]['platform_low']
                    modality['is_cross_mid'] = self.modalityList[-1:][0]['is_cross_mid']
                    if stock.high[0] > self.modalityList[-1:][0]['platform_high']:
                        modality['platform_high'] = stock.high[0]
                        if self.modalityList[-1:][0]['platform_low'] is not None:
                            modality['platform_low'] = self.modalityList[-1:][0]['platform_low']
                    else:
                        # 判断回调后是否创出阶段新低
                        if self.modalityList[-1:][0]['platform_low'] is None:
                            modality['platform_low'] = stock.low[0]
                        else:
                            if stock.low[0] < self.modalityList[-1:][0]['platform_low']:
                                modality['platform_low'] = stock.low[0]
                            else:
                                modality['platform_low'] = self.modalityList[-1:][0]['platform_low']
                    # 做多介入条件：向上突破平台最高价、布林带中轨和长期EMA向上，出现明显的回调
                    cross_price = self.modalityList[-1:][0]['close_cross_mid']
                    platform_high = self.modalityList[-1:][0]['platform_high']
                    enterShort = stock.close[0] < (cross_price + (platform_high - cross_price) * 3 / 4)
                    enterLong = stock.close[0] > platform_high * 1.004 and \
                                self.modalityList[-1:][0]['is_mid_increase'] and \
                                self.modalityList[-1:][0]['is_ema_increase'] and \
                                self.modalityList[-1:][0]['platform_low'] is not None
                    if enterLong:
                        info = self.broker.getcommissioninfo(stock)
                        symbol_multi = info.p.mult
                        total_value = self.broker.getvalue()
                        lots = math.floor(total_value / (symbol_multi * stock.close[0]) / 100) * 100
                        self.loss = self.modalityList[-1:][0]['platform_low']
                        modality['initial_loss'] = self.modalityList[-1:][0]['platform_low']
                        modality['break_high_enter'] = True
                        modality['status'] = 'Signal'
                        self.modalityList.append(modality)
                        # print(pd.DataFrame(self.modalityList))
                        # self.buy_bracket(stock, size=lots, stopprice=self.loss)
                        self.buy(stock, size=lots)
                        self.log(
                            '买入 信号, 交易单价: %.2f, 交易数量: %i, 当前仓位: %i, 初始止损价：%.2f' %
                            (stock.close[0], lots, self.getposition(stock).size, self.loss))
                    elif enterShort:
                        modality['break_low_enter'] = True
                        modality['status'] = 'GiveUp'
                        self.modalityList.append(modality)
                        # print(pd.DataFrame(self.modalityList))
                        _lists = pd.DataFrame(self.modalityList)
                        self.modalitySign = pd.concat([self.modalitySign, _lists], ignore_index=True)
                        self.modalityList = []
                        self.watchSign = False
                    else:
                        modality['status'] = 'Watch'
                        self.modalityList.append(modality)
            else:
                # 已经创建订单
                modality['close_cross_mid'] = self.modalityList[-1:][0]['close_cross_mid']
                modality['platform_high'] = self.modalityList[-1:][0]['platform_high']
                modality['platform_low'] = self.modalityList[-1:][0]['platform_low']
                modality['is_cross_mid'] = self.modalityList[-1:][0]['is_cross_mid']
                loss = stock.low[0] - indicators['atr'] * 1.5
                if self.loss < loss and stock.close[0] > self.modalityList[-1:][0]['platform_high'] * 1.01:
                    # 向上调整移动止损位
                    self.loss = loss
                    modality['moving_lost'] = loss
                    self.log('当前收盘价：%.2f，买入价格：%2f，进行【调整移动止损位】上调至: %.2f' % (
                        stock.close[0], self.buyprice, loss))
                    self.modalityList.append(modality)
                if stock.close[0] < self.loss:
                    # 平仓
                    if self.close(stock):
                        if self.modalityList[-1:][0]['initial_loss'] == self.loss:
                            modality['trigger_loss'] = (stock.close[0], stock.datetime.datetime(0))
                            modality['status'] = 'InitialLoss'
                            self.log('初始止损 信号 (平仓), 交易单价: %.2f, 交易数量: %i' % (
                                stock.close[0], self.getposition(stock).size))
                        else:
                            modality['trigger_atr'] = (stock.close[0], stock.datetime.datetime(0))
                            modality['status'] = 'MovingLoss'
                            self.log('移动止盈 信号 (平仓), 交易单价: %.2f, 交易数量: %i' % (
                            stock.close[0], self.getposition(stock).size))
                        self.modalityList.append(modality)
                        print(pd.DataFrame(self.modalityList))
                else:
                    modality['status'] = 'Wait'
                    self.modalityList.append(modality)

    def shop(self):
        print('Stop')
        # 显示所有列
        pd.set_option('display.max_columns', None)
        # 显示所有行
        pd.set_option('display.max_rows', None)
        print(self.modalitySign)

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.loss = self.datas[0].low[0] - self.indicators[0]['atr'][0] * 1.5
                self.log(
                    '买入 执行, 交易单价: %.2f, 交易数量: %i, 初始止损价：%.2f，当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                    (order.executed.price,
                     order.executed.size,
                     self.loss,
                     self.broker.get_cash(),
                     order.executed.value,
                     order.executed.comm))
                self.buyprice = order.executed.price
                self.orderSign = True
                self.marketposition = 1
            else:  # Sell
                self.log('卖出 执行, 交易单价: %.2f, 交易数量: %i, 当前账户现金余额: %.2f, 交易金额: %.2f，佣金：%.2f' %
                         (order.executed.price,
                          order.executed.size,
                          self.broker.get_cash(),
                          order.executed.value,
                          order.executed.comm))
                _lists = pd.DataFrame(self.modalityList)
                self.modalitySign = pd.concat([self.modalitySign, _lists], ignore_index=True)
                self.modalityList.clear()
                self.watchSign = False
                self.orderSign = False
                self.sellprice = order.executed.price
                self.marketposition = 0
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

    def notify_trade(self, trade):
        if not trade.isclosed:
            return
        self.log('交易利润, 总收入: %.2f, 净利润: %.2f' % (trade.pnl, trade.pnlcomm))
        print()


class PandasStockData(bt.feeds.PandasData):
    lines = ('amount', 'ratio')  # 要添加的线
    params = (
        ('nullvalue', 0.0),
        ('amount', -1),
        ('ratio', -1),
    )


if __name__ == '__main__':
    setting = units.XML2ObjFromFile('./setting.xml').setting
    tdx = TdxSource.Security()
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TradeStrategy)  # 增加一个策略
    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_slippage_fixed(fixed=0.001)  # 设置固定滑点
    cerebro.broker.setcommission(commission=0.00005)  # 设置 0.00005 = 0.005% 的手续费
    conn = units.dbConnect(setting.pgsql)
    stockSecurityData = dict()
    try:
        cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cursor.execute("SELECT * FROM trading.owner_stock_view  Limit 1")
        row = cursor.fetchone()
    except Exception as e:
        logging.warning(e)
    finally:
        units.dbDisconnect(conn)
    # stock_data_15m = StockPandasData(dataname=tdx.get_stock_data(row, '15m').data, timeframe=bt.TimeFrame.Minutes, compression=15)
    stock_data_60m = PandasStockData(dataname=tdx.get_stock_data(row, '60m').data, timeframe=bt.TimeFrame.Minutes,
                                     compression=60)
    # cerebro.adddata(stock_data_15m, name=row['tdx_code'] + '_15M')
    cerebro.adddata(stock_data_60m, name=row['tdx_code'] + '_60M')
    start_value = cerebro.broker.getvalue()
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    final_value = cerebro.broker.getvalue()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    print('Net Profit: %.2f%%' % ((final_value - start_value) / start_value * 100))
    # cerebro.plot()
